import configparser
import os
import sys
import tkinter as tk
from tkinter import filedialog, scrolledtext

import cv2
import numpy as np

from get_coor import get_box
from get_locs import get_files, merge_intersecting_boxes
import plt_fit as pf
from ui_utils import handle_exceptions, TextRedirector, r2_score, LinearRegression, corrcoef


# 保存 k、b 和 r2 到配置文件
def save_config(theta, r2):
    config = configparser.ConfigParser()
    k1, k2, b = theta
    config['PARAMETERS'] = {'k1': str(k1), 'k2': str(k2), 'b': str(b), 'r2': str(r2)}

    with open('AutoFocus.ini', 'w') as configfile:
        config.write(configfile)


# 从配置文件中恢复 k、b 和 r2 的值
def load_config():
    config = configparser.ConfigParser()
    config.read('AutoFocus.ini')

    k1 = float(config['PARAMETERS']['k1'])
    k2 = float(config['PARAMETERS']['k2'])
    b = float(config['PARAMETERS']['b'])
    r2 = float(config['PARAMETERS']['r2'])

    return (k1, k2, b), r2


class ModelTrainerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("自动聚焦模型训练")
        default_kwargs = {'padx': '4px', 'pady': '5px'}

        self.dir_path = tk.StringVar()
        # 创建文件路径输入框
        self.file_path_entry = tk.Entry(root, textvariable=self.dir_path, width=60)
        self.file_path_entry.grid(row=0, columnspan=3, **default_kwargs)

        # 创建选择文件按钮
        self.browse_button = tk.Button(root, text="选择文件夹", command=self.browse_file)
        self.browse_button.grid(row=0, column=3, **default_kwargs)

        # 创建训练模型按钮
        self.train_button = tk.Button(root, text="训练模型", command=self.train_model)
        self.train_button.grid(row=1, column=0, columnspan=4, **default_kwargs)

        self.train_button = tk.Button(root, text="保存参数", command=self.save_cfg)
        self.train_button.grid(row=2, column=0, columnspan=4, **default_kwargs)

        self.btn_pred = tk.Button(root, text="预测模型", command=self.predict)
        self.btn_pred.grid(row=3, column=0, columnspan=4, **default_kwargs)

        self.output = scrolledtext.ScrolledText(root, width=60, height=30, font=("Arial", 12))
        self.output.grid(row=4, columnspan=4, **default_kwargs)
        # 重定向stdout到ScrolledText小部件
        sys.stdout = TextRedirector(self.output, log2file=True)

    def browse_file(self):
        # 打开文件对话框
        file_path = filedialog.askdirectory()
        # 将文件路径显示在输入框中
        self.dir_path.set(file_path)

    def fit(self, X, y):
        # 计算回归系数
        reg = LinearRegression()
        reg.fit(X, y)
        return reg

    @handle_exceptions
    def train_model(self, arg=None):
        if self.dir_path.get() == '':
            raise ValueError("未选择文件夹！")

        file_dir = self.dir_path.get()
        img_path_list, locs = get_files(file_dir)
        coors = []
        boxs = []
        for i, image_name in enumerate(img_path_list):  # 逐一读取图像
            item = cv2.imread(os.path.join(file_dir, image_name))
            cneter, box, _ = get_box(item)
            coors.append(list(cneter))
            boxs.append(box)
        merge_box, _ = merge_intersecting_boxes(boxs)
        # 使用线性回归拟合数据
        matx = np.array(coors)
        arr_x = matx[:, 0]
        reg = self.fit(matx, locs)
        y_true = np.array(locs)
        y_pred = reg.predict(matx)
        r2 = r2_score(y_true, y_pred)
        # 输出 R^2 值
        corr = corrcoef(matx, y_true)
        print("polyfit R^2 值:", r2, 'corr_x:', corr[0], 'corr_y', corr[1])
        draw_img = cv2.imread(os.path.join(file_dir, img_path_list[0]), cv2.IMREAD_COLOR)
        x, y, w, h = merge_box
        cv2.rectangle(draw_img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        pf.plot_image_and_r2_zzz(draw_img, y_true, y_pred, r2, reg.theta)
        self.theta = reg.theta
        self.r2 = r2

    @handle_exceptions
    def save_cfg(self, arg=None):
        save_config(self.theta, self.r2)
        print(f'参数保存成功!')

    @handle_exceptions
    def predict(self, arg=None):
        file_path = filedialog.askopenfilename(initialdir='.')
        print(file_path)
        pre_val = predict(file_path, self.theta)
        print(pre_val)


def predict(path, theta):
    target = cv2.imread(path, cv2.IMREAD_COLOR)
    my_center, _, _ = get_box(target)
    reg = LinearRegression()
    reg.theta = theta
    x = list(my_center)
    return reg.predict(np.array([x]))[0]


#  pyinstaller -F -c -i .\auto.ico -n AutoFoucus  ./autofocus_ui.py
if __name__ == '__main__':
    # 打印其他参数
    if len(sys.argv) > 1:
        writer = TextRedirector(None, log2file=True)
        path = sys.argv[1]
        (theta), r2 = load_config()
        cur_pos = predict(path, theta)
        writer.write(f'cur_pos:{cur_pos}')
        print(f'cur_pos:{cur_pos}')
    else:
        # 创建 Tkinter 应用程序
        root = tk.Tk()
        app = ModelTrainerApp(root)
        root.mainloop()
